#pragma once

#include <arpa/inet.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <sys/socket.h>
#include <time.h>
#include <unistd.h>
#include <cstring>
#include <string>

#include <doca_argp.h>
#include <doca_buf.h>
#include <doca_buf_array.h>
#include <doca_dev.h>
#include <doca_dma.h>
#include <doca_error.h>
#include <doca_log.h>
#include <doca_mmap.h>
#include <doca_pe.h>
#include <doca_ctx.h>
#include <doca_buf_inventory.h>

#include <butil/logging.h>

struct dma_handler { // common handler, some handler to save, every DmaUnit has one
    doca_dev * dev;
    doca_dma * dma;
    doca_ctx * ctx;
    doca_pe * pe;
};

struct dma_config { // common config, device address save
    std::string nic_pcie_addr;
    std::string gpu_pcie_addr; 
    std::string dpu_ip_addr;
    uint32_t task_num;
    doca_dma_task_memcpy_completion_cb_t complet_cb;
    doca_dma_task_memcpy_completion_cb_t error_cb;
    doca_data ctx_user_data;
};

struct dma_memory { // memory config, save pointers
    void * ptr = nullptr; // malloc ptr
    size_t nbytes; // total_len, task_num * task_bytes
    void *export_desc; // current mmap's export pointer
    size_t export_desc_len;
    std::string enemy_export_desc; // enemy_string info
    struct doca_mmap *mmap = nullptr;
    doca_dev * dev;
};

doca_error_t Dma_create(struct dma_handler * dma, struct dma_config * cfg, int mmap_type);

inline doca_error_t open_doca_device_with_pci(const char *pci_addr,
                                              struct doca_dev **retval) {
  struct doca_devinfo **dev_list;
  uint32_t nb_devs;
  uint8_t is_addr_equal = 0;
  doca_error_t res;
  size_t i;

  /* Set default return value */
  *retval = NULL;

  res = doca_devinfo_create_list(&dev_list, &nb_devs);
  if (res != DOCA_SUCCESS) {
    LOG(ERROR) << "Failed to load doca devices list: "
               << doca_error_get_descr(res);
    return res;
  }

  /* Search */
  for (i = 0; i < nb_devs; i++) {
    res = doca_devinfo_is_equal_pci_addr(dev_list[i], pci_addr, &is_addr_equal);
    if (res == DOCA_SUCCESS && is_addr_equal) {
      /* if device can be opened */
      res = doca_dev_open(dev_list[i], retval);
      if (res == DOCA_SUCCESS) {
        doca_devinfo_destroy_list(dev_list);
        return res;
      }
    }
  }

  res = DOCA_ERROR_NOT_FOUND;

  doca_devinfo_destroy_list(dev_list);
  return res;
}

inline doca_error_t init_doca_device(char *nic_pcie_addr,
                                     struct doca_dev **ddev) {
  doca_error_t status;

  if (nic_pcie_addr == NULL || ddev == NULL)
    return DOCA_ERROR_INVALID_VALUE;

  status = open_doca_device_with_pci(nic_pcie_addr, ddev);
  if (status != DOCA_SUCCESS) {
    LOG(ERROR) << "Failed to open DOCA device based on NIC PCI address";
    return status;
  }

  return DOCA_SUCCESS;
}